import tensorflow as tf
import collections
import os

# return list of words from an article
def _read_words(filename):
    with tf.gfile.GFile(filename, "r") as f:
        return f.read().replace("\n", "<eos>").split()

# return dict of {word: id} for words in an article
def _build_vocab(filename):
    # read word as list
    data = _read_words(filename)

    # count word frequency
    counter = collections.Counter(data)
    # return a new list
    # with decreasing frequency, lexicographically ascending order
    # as form of (word, frequency)
    count_pairs = sorted(counter.items(), key=lambda x: (-x[1], x[0]))

    # separate word from its frequency
    words, _ = list(zip(*count_pairs))
    # return a dict of {word: id}
    # where smaller id means high frequency
    word_to_id = dict(zip(words, range(len(words))))

    return word_to_id

# return list of words' ids as order in article if exist in dict
def _file_to_word_ids(filename, word_to_id):
    data = _read_words(filename)
    return [word_to_id[word] for word in data if word in word_to_id]


def raw_data(data_path=None):
    """Load PTB raw data from data directory "data_path".
    Reads PTB text files, converts strings to integer ids,
    and performs mini-batching of the inputs.
    The PTB dataset comes from Tomas Mikolov's webpage:
    http://www.fit.vutbr.cz/~imikolov/rnnlm/simple-examples.tgz
    Args:
        data_path: string path to the directory where simple-examples.tgz has
            been extracted.
    Returns:
        tuple (train_data, valid_data, test_data, vocabulary)
        where each of the data objects can be passed to PTBIterator.
    """

    train_path = os.path.join(data_path, "ptb.train.txt")
    valid_path = os.path.join(data_path, "ptb.valid.txt")
    test_path = os.path.join(data_path, "ptb.test.txt")

    # use train date to build vocab dict
    word_to_id = _build_vocab(train_path)
    # use vocab dict to id words in file
    train_data = _file_to_word_ids(train_path, word_to_id)
    valid_data = _file_to_word_ids(valid_path, word_to_id)
    test_data = _file_to_word_ids(test_path, word_to_id)
    vocabulary = len(word_to_id)
    return train_data, valid_data, test_data, vocabulary
